# 项目模型来自hugging face镜像网站，HF Mirror
# 中译文模型：https://hf-mirror.com/Helsinki-NLP/opus-mt-zh-en/tree/main
# 英译中模型：https://hf-mirror.com/Helsinki-NLP/opus-mt-en-zh/tree/main
import os
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import pipeline
from my_util import Logger

loger = Logger()

def init_text_translate_models():

    try:
        # 加载中译英模型
        model_cn2en_path = os.path.join(os.getcwd(), 'models', 'ch-en')
        # 创建tokenizer
        tokenizer = AutoTokenizer.from_pretrained(model_cn2en_path)
        # 创建模型
        model = AutoModelForSeq2SeqLM.from_pretrained(model_cn2en_path)
        # 创建pipeline
        global pipeline_ch2en
        pipeline_ch2en = pipeline("translation", model=model, tokenizer=tokenizer)

        # 加载英译中模型
        model_en2cn_path = os.path.join(os.getcwd(), 'models', 'en-ch')
        # 创建tokenizer
        tokenizer = AutoTokenizer.from_pretrained(model_en2cn_path)
        # 创建模型
        model = AutoModelForSeq2SeqLM.from_pretrained(model_en2cn_path)
        # 创建pipeline
        global pipeline_en2ch
        pipeline_en2ch = pipeline("translation", model=model, tokenizer=tokenizer)
    except Exception as e:
        # 捕获所有异常，并打印错误信息
        loger.error(f"An error occurred: {e}")
    finally:
        loger.info(f"load text translate models success")
        return


def translate_ch2en(sentence):
    english_res = "unknown"
    try:
        result = pipeline_ch2en(sentence)
        english_res = result[0]['translation_text']
    except Exception as e:
        # 捕获所有异常，并打印错误信息
        loger.error(f"An error occurred: {e}")
    finally:
        loger.info(f"translate {sentence} to {english_res}")
        return english_res

def translate_en2ch(sentence):
    chinese_res = "未知"
    try:
        result = pipeline_en2ch(sentence)
        chinese_res = result[0]['translation_text']
    except Exception as e:
        # 捕获所有异常，并打印错误信息
        loger.error(f"An error occurred: {e}")
    finally:
        loger.info(f"translate {sentence} to {chinese_res}")
        return chinese_res

# if __name__ == "__main__":
#     init_translate_model()
#     print("initializing translation models final")
#     chinese = """
#     六岁时，我家在荷兰的莱斯韦克，房子的前面有一片荒地，
#     我称其为“那地方”，一个神秘的所在，那里深深的草木如今只到我的腰际，
#     当年却像是一片丛林，即便现在我还记得：“那地方”危机四伏，
#     洒满了我的恐惧和幻想。
#     """
#     result = pipeline_ch2en(chinese)
#     english = result[0]['translation_text']
#     print(english)
#
#     result = pipeline_en2ch(english)
#     print(result[0]['translation_text'])